# 提取一副彩色图像中红色，用HIS模型处理，RGB模型对比显示；

import cv2
import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


# bgr_2_hsi、color_slicing来自https://blog.csdn.net/qq_29598161/article/details/102655763
# 参考了hsv提取彩色的方法来自https://blog.csdn.net/gaoyu1253401563/article/details/85253511?utm_medium=distribute.pc_relevant.none-task-blog-2~default~baidujs_title~default-5.no_search_link&spm=1001.2101.3001.4242.4
# HSV提取指定其他所有颜色的方法https://blog.csdn.net/qq_41895190/article/details/82791426
def bgr_2_hsi(image):
    out = np.copy(image)
    out_slicing = np.zeros(image.shape, np.uint8)

    for x in range(image.shape[0]):
        # print(str(int(x / image.shape[0] * 100)) + "%")
        for y in range(image.shape[1]):
            b, g, r = image[x][y]
            b, g, r = int(b), int(g), int(r)
            i_s = np.sum([b, g, r])
            i = i_s / 3

            # i == 0, s and h is no sense
            if i_s == 0:
                i = 0
                s = 0
                h = 0
                out[x][y] = h, s, i
                continue

            s = (1 - (3 * np.min([b, g, r])) / i_s) * 255

            # s == 0 h is no sense
            if s == 0:
                h = 0
                out[x][y] = h, s, i
                continue

            thea = np.arccos((2 * r - g - b) / (2 * np.sqrt((r - g) ** 2 + (r - b) * (g - b))))
            if g >= b:
                h1 = thea
            else:
                h1 = np.pi * 2 - thea
            h1 = np.rad2deg(h1)
            # slicing
            if (int(h1) in range(0, 11) or int(h1) in range(350, 361)) and s / 255 > 0.1:
                print(int(h1))
                out_slicing[x][y] = image[x][y]

            h = h1 / 360 * 255
            out[x][y] = h, s, i

    return out


# 读入的图像是BGR空间图像
frame = cv2.imread("img/color2.png")

# frame_to_hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# 部分1：将BGR空间的图片转换到HSI空间
frame_to_hsi = bgr_2_hsi(frame)

# 部分2：
# 每个分量设定两个阈值
#    H_min < H < H_max  AND
#    S_min < S < S_max  AND
#     I_min < I < I_max
# 定义红色
lower_red = np.array([0, 100, 100])
upper_red = np.array([10, 255, 255])

# 部分3：
# 从HSI图像中截取出蓝色、绿色、红色，即获得相应的掩膜
# cv2.inRange()函数是设置阈值去除背景部分，得到想要的区域

red_mask = cv2.inRange(frame_to_hsi, lower_red, upper_red)

# 部分4：
# 将原图像和mask(掩膜)进行按位与

red_res = cv2.bitwise_and(frame, frame, mask=red_mask)

# 最后得到要分离出的颜色图像

# 部分5:将BGR空间下的图片转换成RGB空间下的图片
frame = frame[:, :, ::-1]
red_res = red_res[:, :, ::-1]

# 部分6：显示图像

plt.subplot(131), plt.imshow(frame) ,plt.title("原图")
plt.subplot(132), plt.imshow(red_mask), plt.title("红色遮罩")
plt.subplot(133), plt.imshow(red_res), plt.title("红色")

plt.show()


def color_slicing(image, center, w):
    """
    :param image:
    :param center: b, g, r ib range 0 ~ 255
    :param w: width
    :return:
    """
    out = np.zeros(image.shape, np.uint8)
    for x in range(image.shape[0]):
        for y in range(image.shape[1]):
            r_b, r_g, r_r = center
            a_b, a_g, a_r = image[x][y]
            if abs(r_b - a_b) < w / 2 and abs(r_g - a_g) < w / 2 and abs(r_r - a_r) < w / 2:
                out[x][y] = image[x][y]
    return out


img = cv2.imread("img/color2.png")
img_to_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_slice_red = color_slicing(img_to_rgb, (241, 68, 58), 0.2549 * 255)
plt.subplot(121), plt.imshow(img_to_rgb), plt.title("原图RGB")
plt.subplot(122), plt.imshow(img_slice_red), plt.title("红色")
plt.show()
